import numpy as np
from tensorflow.keras.models import load_model

# 加载模型
model = load_model('D:\Project\Pycharm Project\pythonProject4\models\snake_model.h5')

# 游戏参数
WIDTH, HEIGHT = 10, 10  # 游戏区域大小
SNAKE_INIT_LENGTH = 3  # 蛇的初始长度
SNAKE_SPEED = 1  # 蛇的移动速度

# 方向定义
UP = 0
RIGHT = 3
DOWN = 1
LEFT = 2

# 初始化游戏状态
snake = [(WIDTH // 2, HEIGHT // 2)] * SNAKE_INIT_LENGTH
direction = RIGHT
food = (np.random.randint(0, WIDTH), np.random.randint(0, HEIGHT))
score = 0

# 游戏循环
while True:
    # 获取当前状态
    state = np.zeros((1, WIDTH * HEIGHT), dtype=np.float32)
    for x, y in snake:
        state[0, x * WIDTH + y] = 1
    state[0, food[0] * WIDTH + food[1]] = 2

    # 使用模型预测下一步动作
    prediction = model.predict(state)
    action = np.argmax(prediction[0])
    if action == UP:
        print("模型预测的方向为上")
    elif action == RIGHT:
        print("模型预测的方向为右")
    elif action == DOWN:
        print("模型预测的方向为下")
    elif action == LEFT:
        print("模型预测的方向为左")

    # 询问用户是否要覆盖模型的预测
    user_override = input("是否要覆盖模型的预测？(y/n): ")
    if user_override.lower() == 'y':
        # 用户输入
        action = int(input("请输入方向（0-3）:0-上,1-右,2-下,3-左"))

    # 更新方向
    if action == UP and direction != DOWN:
        direction = UP
    elif action == RIGHT and direction != LEFT:
        direction = RIGHT
    elif action == DOWN and direction != UP:
        direction = DOWN
    elif action == LEFT and direction != RIGHT:
        direction = LEFT

    # 移动蛇
    head = snake[0]
    if direction == UP:
        new_head = (head[0], head[1] - SNAKE_SPEED)
    elif direction == RIGHT:
        new_head = (head[0] + SNAKE_SPEED, head[1])
    elif direction == DOWN:
        new_head = (head[0], head[1] + SNAKE_SPEED)
    elif direction == LEFT:
        new_head = (head[0] - SNAKE_SPEED, head[1])

    # 检查是否吃到食物
    if new_head == food:
        snake.insert(0, new_head)
        score += 1
        food = (np.random.randint(0, WIDTH), np.random.randint(0, HEIGHT))
    else:
        snake.insert(0, new_head)
        snake.pop()

    # 检查是否撞到墙壁或自己
    if (new_head[0] < 0 or new_head[0] >= WIDTH or
            new_head[1] < 0 or new_head[1] >= HEIGHT or
            new_head in snake[1:]):
        print(f"游戏结束！得分: {score}")
        break

    # 打印游戏状态
    print("-" * (WIDTH + 2))
    for y in range(HEIGHT):
        print("|", end="")
        for x in range(WIDTH):
            if (x, y) in snake:
                print("O", end="")
            elif (x, y) == food:
                print("X", end="")
            else:
                print(" ", end="")
        print("|")
    print("-" * (WIDTH + 2))
    print(f"得分: {score}")
